import os
from os.path import join
import sys
sys.path = ['/home/hzh/.local/lib/python3.5/site-packages/'] + sys.path
import cv2
from datetime import timedelta
import numpy as np
import shutil
from typing import Tuple, Any, List, Dict
# import matplotlib.pyplot as plt

IMG_DIFF_THRES = 0.3
IMG_WHITE_THRES = 100
IMG_EMPTY_THRES = 0.01

cap = cv2.VideoCapture("/home/hzh/Desktop/downloads/Sabaha.mp4")
fps = cap.get(cv2.CAP_PROP_FPS)

timestamps = [cap.get(cv2.CAP_PROP_POS_MSEC)]
# calc_timestamps = [0.0]

start_sec= 75
end_sec = 10000000
step = 1

GAP=int(1000/fps)

def preprocess(img):
    idx = np.logical_or(np.logical_or(img[:,:, 0] < 200, img[:, :, 1] < 200),
                         img[:, :, 2] < 200)
    img[idx]=(0, 0, 0)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    gray[gray>IMG_WHITE_THRES] = 255
    return gray

def slice_bottom(img):
    # print(img.shape)
    b = img[926:]
    return b

def count_white(img):
    c = np.sum(img==255)
    print("white cnt=", c)
    return c

def is_img_diff(a, b):
    # diff = np.sum(a!=b)
    diff_img = a-b
    diff = count_white(diff_img)
    ratio = diff / min(count_white(a), count_white(b))
    print("diff cnt =", diff)
    print("diff ratio=", ratio)
    return ratio > IMG_DIFF_THRES

def is_empty_diff(img):
    r = count_white(img)/(img.shape[0]*img.shape[1])
    print("white ratio=", r)
    return r < IMG_EMPTY_THRES

class DB_Writer:
    def __init__(self, db_dir):
        self._db_dir = db_dir
        if os.path.exists(db_dir):
            cmd = input("%s already exist. Delete? y/[n]"%db_dir)
            if cmd in "yY":
                shutil.rmtree(db_dir)
            else:
                exit(-1)
        os.makedirs(db_dir)
        self._index_f = open(join(db_dir, "index.yaml"), 'w')
        self._cnt = 1
        self._prev = None

    def append(self, img, tm):
        fn = "%d.png" % self._cnt
        cv2.imwrite(join(self._db_dir, fn), img)
        self._index_f.write("%d: %s\n" % (self._cnt, tm))
        self._cnt += 1
        self._prev = img

    def __getitem__(self, i):
        if -1 == i:
            return self._prev
        else:
            raise Exception("only '-1' is support in [] operator of DB_Writer")


class SubtitleSeq(object):
    db = DB_Writer("./bottom_imgs")
    def show(self):
        for s, t, im in zip(self.sbt_seq, self.tm_seq, self.img_seq):
            tm = str(timedelta(seconds=t/1000))
            fp = "./data/sub_imgs/%s.png"%(tm)
            cv2.imwrite(fp, im)
            print(t/1000)
            print(s)

    def add(self, img, tm):
        if is_empty_diff(img):
            # cv2.imshow("empty", np.vstack(img))
            # cv2.waitKey(1)
            pass
        else:
            if None is self.db[-1]:
                self.db.append(img, tm)
            elif is_img_diff(self.db[-1], img):
                print("+", tm)
                # cv2.imshow("diff", np.vstack([self.db[-1], img]))
                # cv2.waitKey(1)
                self.db.append(img, tm)
            else:
                # cv2.imshow("same", np.vstack([self.db[-1], img]))
                # cv2.waitKey(1)
                pass


if "__main__" == __name__:
    sbt = SubtitleSeq()
    while(cap.isOpened()):
        frame_exists, curr_frame = cap.read()
        if not frame_exists:
            print("no frame")
            break
        tm = cap.get(cv2.CAP_PROP_POS_MSEC)
        # print(tm)
        if tm > start_sec*1000:
            start_sec += step
            # cv2.imshow("movie", curr_frame)
            bt_img = slice_bottom(curr_frame)
            bt_img = preprocess(bt_img)
            # cv2.imshow("bottom", bt_img)
            # cv2.imwrite("subtitle.png", bt_img)
            # cv2.waitKey(GAP)
            #try:
            #    s = extract_chinese(bt_img)
            #except Exception as e:
            #    print "exception:", e
            #print(s)
            sbt.add(bt_img, tm)
            # timestamps.append(cap.get(cv2.CAP_PROP_POS_MSEC))
            # calc_timestamps.append(calc_timestamps[-1] + 1000/fps)
        if tm > end_sec*1000:
            print("time ends")
            break

    cap.release()
